import os
import sys

def to_head(projectpath):
    pathlayers = os.path.join(projectpath, 'layers/').replace('\\', '/')
    return r"""
\documentclass[border=8pt, multi, tikz]{standalone} 
\usepackage{import}
\subimport{""" + pathlayers + r"""}{init}
\usetikzlibrary{positioning}
\usetikzlibrary{3d} %for including external image 
"""


def to_cor():
    return r"""
\def\ConvColor{rgb:yellow,5;red,2.5;white,5}
\def\ConvReluColor{rgb:yellow,5;red,5;white,5}
\def\PoolColor{rgb:red,1;black,0.3}
\def\UnpoolColor{rgb:blue,2;green,1;black,0.3}
\def\FcColor{rgb:blue,5;red,2.5;white,5}
\def\FcReluColor{rgb:blue,5;red,5;white,4}
\def\SoftmaxColor{rgb:magenta,5;black,7}   
\def\SumColor{rgb:blue,5;green,15}
"""


def to_begin():
    return r"""
\newcommand{\copymidarrow}{\tikz \draw[-Stealth,line width=0.8mm,draw={rgb:blue,4;red,1;green,1;black,3}] (-0.3,0) -- ++(0.3,0);}

\begin{document}
\begin{tikzpicture}
\tikzstyle{connection}=[ultra thick,every node/.style={sloped,allow upside down},draw=\edgecolor,opacity=0.7]
\tikzstyle{copyconnection}=[ultra thick,every node/.style={sloped,allow upside down},draw={rgb:blue,4;red,1;green,1;black,3},opacity=0.7]
"""


# layers definition

def to_input(pathfile, to='(-3,0,0)', width=8, height=8, name="temp"):
    return r"""
\node[canvas is zy plane at x=0] (""" + name + """) at """ + to + """ {\includegraphics[width=""" + str(
        width) + "cm" + """,height=""" + str(height) + "cm" + """]{""" + pathfile + """}};
"""


# Conv
def to_Conv(name, s_filer=256, n_filer=64, offset="(0,0,0)", to="(0,0,0)", width=1, height=40, depth=40, caption=" "):
    return r"""
\pic[shift={""" + offset + """}] at """ + to + """ 
    {Box={
        name=""" + name + """,
        caption=""" + caption + r""",
        xlabel={{""" + str(n_filer) + """, }},
        zlabel=""" + str(s_filer) + """,
        fill=\ConvColor,
        height=""" + str(height) + """,
        width=""" + str(width) + """,
        depth=""" + str(depth) + """
        }
    };
"""


# Conv,Conv,relu
# Bottleneck
def to_ConvConvRelu(name, s_filer=256, n_filer=(64, 64), offset="(0,0,0)", to="(0,0,0)", width=(2, 2), height=40,
                    depth=40, caption=" "):
    return r"""
\pic[shift={ """ + offset + """ }] at """ + to + """ 
    {RightBandedBox={
        name=""" + name + """,
        caption=""" + caption + """,
        xlabel={{ """ + str(n_filer[0]) + """, """ + str(n_filer[1]) + """ }},
        zlabel=""" + str(s_filer) + """,
        fill=\ConvColor,
        bandfill=\ConvReluColor,
        height=""" + str(height) + """,
        width={ """ + str(width[0]) + """ , """ + str(width[1]) + """ },
        depth=""" + str(depth) + """
        }
    };
"""


# Pool
def to_Pool(name, offset="(0,0,0)", to="(0,0,0)", width=1, height=32, depth=32, opacity=0.5, caption=" "):
    return r"""
\pic[shift={ """ + offset + """ }] at """ + to + """ 
    {Box={
        name=""" + name + """,
        caption=""" + caption + r""",
        fill=\PoolColor,
        opacity=""" + str(opacity) + """,
        height=""" + str(height) + """,
        width=""" + str(width) + """,
        depth=""" + str(depth) + """
        }
    };
"""


# unpool4,
def to_UnPool(name, offset="(0,0,0)", to="(0,0,0)", width=1, height=32, depth=32, opacity=0.5, caption=" "):
    return r"""
\pic[shift={ """ + offset + """ }] at """ + to + """ 
    {Box={
        name=""" + name + r""",
        caption=""" + caption + r""",
        fill=\UnpoolColor,
        opacity=""" + str(opacity) + """,
        height=""" + str(height) + """,
        width=""" + str(width) + """,
        depth=""" + str(depth) + """
        }
    };
"""


def to_ConvRes(name, s_filer=256, n_filer=64, offset="(0,0,0)", to="(0,0,0)", width=6, height=40, depth=40, opacity=0.2,
               caption=" "):
    return r"""
\pic[shift={ """ + offset + """ }] at """ + to + """ 
    {RightBandedBox={
        name=""" + name + """,
        caption=""" + caption + """,
        xlabel={{ """ + str(n_filer) + """, }},
        zlabel=""" + str(s_filer) + r""",
        fill={rgb:white,1;black,3},
        bandfill={rgb:white,1;black,2},
        opacity=""" + str(opacity) + """,
        height=""" + str(height) + """,
        width=""" + str(width) + """,
        depth=""" + str(depth) + """
        }
    };
"""


# ConvSoftMax
def to_ConvSoftMax(name, s_filer=40, offset="(0,0,0)", to="(0,0,0)", width=1, height=40, depth=40, caption=" "):
    return r"""
\pic[shift={""" + offset + """}] at """ + to + """ 
    {Box={
        name=""" + name + """,
        caption=""" + caption + """,
        zlabel=""" + str(s_filer) + """,
        fill=\SoftmaxColor,
        height=""" + str(height) + """,
        width=""" + str(width) + """,
        depth=""" + str(depth) + """
        }
    };
"""


# SoftMax
def to_SoftMax(name, s_filer=10, offset="(0,0,0)", to="(0,0,0)", width=1.5, height=3, depth=25, opacity=0.8,
               caption=" "):
    return r"""
\pic[shift={""" + offset + """}] at """ + to + """ 
    {Box={
        name=""" + name + """,
        caption=""" + caption + """,
        xlabel={{" ","dummy"}},
        zlabel=""" + str(s_filer) + """,
        fill=\SoftmaxColor,
        opacity=""" + str(opacity) + """,
        height=""" + str(height) + """,
        width=""" + str(width) + """,
        depth=""" + str(depth) + """
        }
    };
"""


def to_Sum(name, offset="(0,0,0)", to="(0,0,0)", radius=2.5, opacity=0.6):
    return r"""
\pic[shift={""" + offset + """}] at """ + to + """ 
    {Ball={
        name=""" + name + """,
        fill=\SumColor,
        opacity=""" + str(opacity) + """,
        radius=""" + str(radius) + """,
        logo=$+$
        }
    };
"""


def to_connection(of, to):
    return r"""
\draw [connection]  (""" + of + """-east)    -- node {\midarrow} (""" + to + """-west);
"""


def to_skip(of, to, pos=1.25):
    return r"""
\path (""" + of + """-southeast) -- (""" + of + """-northeast) coordinate[pos=""" + str(pos) + """] (""" + of + """-top) ;
\path (""" + to + """-south)  -- (""" + to + """-north)  coordinate[pos=""" + str(pos) + """] (""" + to + """-top) ;
\draw [copyconnection]  (""" + of + """-northeast)  
-- node {\copymidarrow}(""" + of + """-top)
-- node {\copymidarrow}(""" + to + """-top)
-- node {\copymidarrow} (""" + to + """-north);
"""


def to_end():
    return r"""
\end{tikzpicture}
\end{document}
"""


def to_generate(arch, pathname="file.tex"):
    with open(pathname, "w") as f:
        for c in arch:
            print(c)
            f.write(c)


# defined your arch
arch = [
    to_head('..'),
    to_cor(),
    to_begin(),
    to_input('../../data/train_data/test_1588.jpg', name="input0"),

    # Layer 1
    to_Conv("conv1", 26, 25, offset="(0,0,0)", to="(0,0,0)", height=40, depth=40, width=2, caption="Conv1"),
    to_Pool("pool1", offset="(1.5,0,0)", to="(conv1-east)", height=32, depth=32, width=1, caption="MaxPool1"),  # increased offset
    to_connection("input0", "conv1"),  # input0 to conv1

    # Layer 2
    to_Conv("conv2", 11, 50, offset="(2.0,0,0)", to="(pool1-east)", height=32, depth=32, width=4, caption="Conv2"),  # increased offset
    to_Pool("pool2", offset="(1.5,0,0)", to="(conv2-east)", height=24, depth=24, width=1, caption="MaxPool2"),  # increased offset
    to_connection("pool1", "conv2"),

    # Flatten
    to_Conv("flatten", 1280, 1280, offset="(2.5,0,0)", to="(pool2-east)", height=1, depth=40, width=40, caption="Flatten"),
    to_connection("pool2", "flatten"),

    # FC layers
    to_SoftMax(name='fc1', s_filer=1024, offset="(4,0,0)", to="(flatten-east)", width=1.5, height=1.5, depth=100,
               opacity=0.8, caption='FC1'),
    to_connection("flatten", "fc1"),

    to_SoftMax(name='fc2', s_filer=128, offset="(2,0,0)", to="(fc1-east)", width=1.5, height=1.5, depth=50,
               opacity=0.8, caption='FC2'),
    to_connection("fc1", "fc2"),

    to_SoftMax(name='fc3', s_filer=17, offset="(2,0,0)", to="(fc2-east)", width=1.5, height=1.5, depth=5,
               opacity=0.8, caption='FC3'),
    to_connection("fc2", "fc3"),

    to_end()
]

def main():
    namefile = str(sys.argv[0]).split('.')[0]
    to_generate(arch, namefile + '.tex')

if __name__ == '__main__':
    main()
